import torch
import torch.nn as nn
import torch.nn.functional as F

class LeNet5(nn.Module):
    """
    LeNet-5卷积神经网络模型
    原始结构来自Yann LeCun 1998年的论文
    """
    def __init__(self):
        super(LeNet5, self).__init__()
        # 第一个卷积层 (1x32x32 -> 6x28x28)
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        # 第一个池化层 (6x28x28 -> 6x14x14)
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 第二个卷积层 (6x14x14 -> 16x10x10)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        # 第二个池化层 (16x10x10 -> 16x5x5)
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
        # 第三个卷积层 (16x5x5 -> 120x1x1)
        self.conv3 = nn.Conv2d(16, 120, kernel_size=5)
        # 全连接层1 (120 -> 84)
        self.fc1 = nn.Linear(120, 84)
        # 全连接层2 (84 -> 10)
        self.fc2 = nn.Linear(84, 10)
    
    def forward(self, x):
        # 确保输入是32x32
        if x.size(2) != 32:
            # 进行上采样到32x32
            x = F.interpolate(x, size=(32, 32), mode='bilinear', align_corners=True)
        
        # 第一个卷积+激活+池化
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        
        # 第二个卷积+激活+池化
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        
        # 第三个卷积+激活
        x = F.relu(self.conv3(x))
        # 展平特征图，只保留通道维度 (batch_size, 120, 1, 1) -> (batch_size, 120)
        x = x.view(-1, 120)
        
        # 全连接层1
        x = F.relu(self.fc1(x))
        # 全连接层2（输出层）
        x = self.fc2(x)
        
        return x